package com.zwf.checkstyledemo1.checks;

import com.puppycrawl.tools.checkstyle.api.*;
import java.util.*;

public class NullDereferenceCheck extends AbstractCheck {

    // 存储变量状态 (null, not-null, unknown)
    private final Map<String, VariableState> variableStates = new HashMap<>();
    private final Stack<Set<String>> scopeStack = new Stack<>();

    @Override
    public int[] getDefaultTokens() {
        return new int[] {
                TokenTypes.VARIABLE_DEF,    // 变量定义
                TokenTypes.ASSIGN,           // 赋值
                TokenTypes.METHOD_CALL,      // 方法调用
                TokenTypes.EXPR,             // 表达式
                TokenTypes.LITERAL_IF,       // if语句
                TokenTypes.LITERAL_RETURN,   // return语句
                TokenTypes.LCURLY,           // 代码块开始
                TokenTypes.RCURLY            // 代码块结束
        };
    }

    @Override
    public int[] getAcceptableTokens() {
        return new int[0];
    }

    @Override
    public int[] getRequiredTokens() {
        return new int[0];
    }

    @Override
    public void beginTree(DetailAST rootAST) {
        variableStates.clear();
        scopeStack.clear();
        enterScope();
    }

    @Override
    public void visitToken(DetailAST ast) {
        switch (ast.getType()) {
            case TokenTypes.LCURLY:
                enterScope();
                break;

            case TokenTypes.RCURLY:
                exitScope();
                break;

            case TokenTypes.VARIABLE_DEF:
                handleVariableDefinition(ast);
                break;

            case TokenTypes.ASSIGN:
                handleAssignment(ast);
                break;

            case TokenTypes.METHOD_CALL:
                handleMethodCall(ast);
                break;

            case TokenTypes.EXPR:
                checkDereference(ast);
                break;

            case TokenTypes.LITERAL_IF:
                handleIfStatement(ast);
                break;
        }
    }

    private void enterScope() {
        scopeStack.push(new HashSet<>());
    }

    private void exitScope() {
        if (scopeStack.isEmpty()) {
            return;
        }
        Set<String> currentScopeVars = scopeStack.pop();
        for (String var : currentScopeVars) {
            variableStates.remove(var);
        }
    }

    private void handleVariableDefinition(DetailAST ast) {
        if (scopeStack.isEmpty()) {
            return;
        }
        String varName = ast.findFirstToken(TokenTypes.IDENT).getText();
        scopeStack.peek().add(varName);
        variableStates.put(varName, VariableState.UNKNOWN);
    }

    private void handleAssignment(DetailAST ast) {
        DetailAST target = ast.getFirstChild();
        if (target.getType() == TokenTypes.IDENT) {
            String varName = target.getText();
            DetailAST expr = ast.getFirstChild().getNextSibling();

            if (isNullLiteral(expr)) {
                variableStates.put(varName, VariableState.NULL);
            } else if (isNotNullExpression(expr)) {
                variableStates.put(varName, VariableState.NOT_NULL);
            } else {
                variableStates.put(varName, VariableState.UNKNOWN);
            }
        }
    }

    private void handleMethodCall(DetailAST ast) {
        // 检查可能返回 null 的方法调用
        if (isMethodThatMayReturnNull(ast)) {
            DetailAST parent = ast.getParent();
            if (parent.getType() == TokenTypes.EXPR) {
                // 忽略返回值
            } else {
                log(ast, "方法返回可能为null", ast.getText());
            }
        }
    }

    private void checkDereference(DetailAST expr) {
        DetailAST child = expr.getFirstChild();
        if (child != null) {
            if (child.getType() == TokenTypes.DOT) {
                DetailAST firstOperand = child.getFirstChild();
                if (firstOperand.getType() == TokenTypes.IDENT) {
                    String varName = firstOperand.getText();
                    VariableState state = variableStates.getOrDefault(varName, VariableState.UNKNOWN);

                    if (state == VariableState.NULL) {
                        log(expr, "空指针风险", varName);
                    }
                }
            }
        }
    }

    private void traverseBlock(DetailAST block) {
        DetailAST child = block.getFirstChild();
        while (child != null) {
            visitToken(child);
            if (child.hasChildren()) {
                traverseBlock(child);
            }

            child = child.getNextSibling();
        }
    }

    private void handleIfStatement(DetailAST ifAst) {
        DetailAST condition = ifAst.findFirstToken(TokenTypes.LPAREN).getNextSibling();
        if (isNullCheckCondition(condition)) {
            // 提取被检查的变量名
            String checkedVar = extractCheckedVariable(condition);

            // 在 if 块内，变量状态已知
            if (checkedVar != null) {
                DetailAST ifBlock = ifAst.findFirstToken(TokenTypes.SLIST);
                if (ifBlock != null) {
                    // 在 if 块内标记变量状态
                    variableStates.put(checkedVar, VariableState.NOT_NULL);
                    // 递归处理 if 块内的语句
                    traverseBlock(ifBlock);
                    // 恢复状态
                    variableStates.put(checkedVar, VariableState.UNKNOWN);
                }
            }
        }
    }

    private boolean isNullCheckCondition(DetailAST condition) {
        // 检查条件是否为 "var == null" 或 "var != null"
        if (condition.getType() == TokenTypes.EQUAL ||
                condition.getType() == TokenTypes.NOT_EQUAL) {

            DetailAST first = condition.getFirstChild();
            DetailAST second = first.getNextSibling();

            return (isNullLiteral(first) && second.getType() == TokenTypes.IDENT) ||
                    (isNullLiteral(second) && first.getType() == TokenTypes.IDENT);
        }
        return false;
    }

    private String extractCheckedVariable(DetailAST condition) {
        DetailAST first = condition.getFirstChild();
        DetailAST second = first.getNextSibling();

        if (first.getType() == TokenTypes.IDENT) return first.getText();
        if (second.getType() == TokenTypes.IDENT) return second.getText();
        return null;
    }

    private boolean isNullLiteral(DetailAST expr) {
        return expr != null && expr.getType() == TokenTypes.LITERAL_NULL;
    }

    private boolean isNotNullExpression(DetailAST expr) {
        // 检查表达式是否为 "new ..." 或字面量
        return expr != null &&
                (expr.getType() == TokenTypes.LITERAL_NEW ||
                        expr.getType() == TokenTypes.STRING_LITERAL ||
                        expr.getType() == TokenTypes.NUM_INT ||
                        expr.getType() == TokenTypes.NUM_LONG ||
                        expr.getType() == TokenTypes.NUM_FLOAT ||
                        expr.getType() == TokenTypes.NUM_DOUBLE);
    }

    private boolean isMethodThatMayReturnNull(DetailAST methodCall) {
        // 实际项目中应使用更复杂的启发式方法
        String methodName = methodCall.getFirstChild().getText();

        // 可能返回 null 的方法名模式
        return methodName.startsWith("get") ||
                methodName.startsWith("find") ||
                methodName.startsWith("lookup") ||
                methodName.startsWith("create") ||
                methodName.endsWith("OrNull");
    }

    private enum VariableState {
        NULL, NOT_NULL, UNKNOWN
    }
}
